!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
! On the Optimal Design of Transfers and Income-Tax Progressivity
! Optimal transition from calibrated economy to steady state with affine tax
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

program Main_PreTesting

    use omp_lib
    use Globals
    use Toolbox
    use ModuleSolveEqm
    use ModuleQTRANSITION

    implicit none 
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Local variable declaration
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    
    ! For submission of many programs
    character(len=5)    :: seq_no_s     ! read input from Linux directly
    integer             :: seq_no       ! sequence number of program 
    integer             :: no_assign    ! number of steady states assigned to each program
    integer             :: loop_start   ! index of first steady state computed by this program
    integer             :: loop_end     ! index of last steady state computed by this program
    logical             :: exist        ! to check whether file already exists

    ! For optimization
    integer, parameter :: N = 1                 ! number of parameters to be optimized
    integer, parameter  :: nlambda = 61         ! number of points for tax level
    integer, parameter :: II = nlambda          ! number of grid points
    integer, parameter  :: no_prog = 61         ! number of programs to be run (II needs to be multiple of no_prog)
    integer :: ia, ian                          ! counter
    real(8) :: paramSOB(II, N)                  ! parameter matrix
    real(8) :: fvalSOB(II)                      ! welfare at the equilibria 
    real(8) :: fvalSOB_ss(II)                   ! welfare only for steady states
    integer, parameter :: N_info = 16           ! number of eqm results to be stored
    integer ::  isave                           ! counter for saving
    real(8) :: info(II,N_info)                  ! matrix of results over grid
    real(8) :: lambda_grid(nlambda)             ! grid for tax level
    integer :: ibn, icn, idn, ix                ! counter
    real(8) :: welfare                          ! welfare associated with equilibrium
    real(8) :: param_vec_init(3), param_vec_lr(3)   ! vector of tax function parameters: input for function solving for equilibrium given taxes
    
    ! Errors 
    real(8) ::  errors_vec_Sob(II,4)             ! errors for each Sobol point
    real(8) ::  errors_vec_trans(II,2,T_TR)      ! errors for transition
    real(8) :: error_scalar_trans(II,2)          ! maxval of errors_vec_trans in transition
    
    ! Calibrated parameters
    real(8) :: ParamCalib(5), eqm_results(5)    ! calibrated parameters
    
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Read sequence number of this program
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    
    if ( COMMAND_ARGUMENT_COUNT() /= 1) then                           
        print *, 'ERROR: A command line argument specifying the parameter is required'
        STOP
    endif 
    print *, 'about to read'
    call get_command_argument(1, seq_no_s)
    print *, 'done reading'
    read(seq_no_s,*) seq_no              
    print *, 'done writing seq_no'            
    print *, 'seq_no    = ', seq_no
    print *, 'seq_no_s   = ', trim(seq_no_s)
    
    
    ! Initialize info file storing results
    if (seq_no == 1) then
                
        info=0
        open(1,  file = 'Output/info.txt',      status = 'unknown')
        write(1, '(*(f18.12))') ( Info(1,isave), isave = 1, N_info )
        close(1)
    
    endif

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Assign calibrated parameters
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    
    ! Choose normalization version
    norm_version = 1            ! indicator for choice of normalization: 1) mean income; 2) median income; 3) mean income of the third quintile
    
    ! Calibrated parameters
    open(11,  file = 'Input/ParamCalib.txt',      status = 'unknown')
    read(11, *) ParamCalib
    close(11)
    B           = ParamCalib(2)             ! labor disutility parameter
    beta        = ParamCalib(1)             ! discount factor
    G           = ParamCalib(3)             ! government spending
    Debt        = ParamCalib(4)             ! government debt
    ymean       = ParamCalib(5)             ! mean income
    ymeanCalib  = ParamCalib(5)             ! mean income
    
    ! Set tax function parameters
    open(11,  file = 'Input/eqm_results.txt',      status = 'unknown')
    read(11, *) eqm_results
    close(11)
    param_vec_init(1) = eqm_results(3)     ! gamma
    param_vec_init(2) = eqm_results(2)     ! lambda
    param_vec_init(3) = eqm_results(5)     ! rt
    omegat = 0d0
    print *, 'Tax parameters initially: gamma, lambda, rt = ', param_vec_init
    print *, ''

    ! Guesses for equilibrium objects
    mt     = eqm_results(4)             ! mt
    r      = eqm_results(1)             ! interest rate
    lambda = param_vec_init(2)          ! lambda
    print*, 'Guesses for r = ', r,' and mt = ', mt
    print *, ''
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Compute initial steady state
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    ! Solve for initial steady state
    print *, 'Computing initial steady-state'
    print *, ''
    
    ! Solve for equilibrium given tax function parameters
    welfare = welfare_eqm(param_vec_init,1)
    
    ! Store eqm objects
    r_init = r                  ! initial equilibrium interest rate
    mt_init = mt                ! initial equilibrium level transfer
    error_vec_init = error_vec  ! initial equilibrium vector of errors
    Kagg_init = Kagg            ! initial equilibrium capital stock

    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Generate grids for grid search over tax systems
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    
    ! Grids for individual tax-transfer parameters
    lambda_grid = LINSPACE_FET(0.35d0, 0.65D0,nlambda)    ! tax level        
    
    ! Arrange into paramSOB matrix
    ix=0
    do ian = 1, nlambda
        ix=ix+1
        paramSOB(ix,1) = lambda_grid(ian)
    enddo
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Compute transitions for points assigned to this program
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    
    ! How many steady states for each program?
    no_assign = II/no_prog
    
    ! Compute the equilibrium for each point assigned to this program
    loop_start = (seq_no-1)*no_assign+1
    loop_end = (seq_no-1)*no_assign+no_assign
    
    
    do ia = loop_start, loop_end              
       
        ! Show iteration number
        print*, '** iter=', ia
        
        ! Tax-transfer parameter for final steady state
        param_vec_lr(1) = 0d0               ! tax progressivity
        param_vec_lr(2) = paramSOB(ia,1)    ! tax level
        param_vec_lr(3) = 0d0               ! phase-out
        
        ! Guesses for eqm objeccts
        r       = 0.02d0
        mt      = 0.4d0
        
        ! Compute final steady state
        fvalSOB_ss(ia) = welfare_eqm(param_vec_lr,2)
        
        ! Store eqm objects
        r_lr = r                    ! new equilibrium interest rate
        lambda_lr = lambda          ! new equilibrium tax function level
        mt_lr = mt                  ! new equilibrium level transfer
        error_vec_lr  = error_vec   ! new equilibrium vector of errors
     
        ! Show result of this steady state
        print *, ''
        print *, '** done with the steady state with lambda = ', lambda
        print *, '** Output r, mt=', r, mt
        
        ! If steady state converged: transition
        ! 1: asset market; 2: govt BC; 3: VFI; 4: measure
        if (error_vec(1) < 1d-5 .AND. error_vec(2) < 1d-5 .AND. error_vec(3) < 5d-10 .AND. error_vec(4) < 1d-11) then
            
            ! Solving for transition
            print *, ''
            print *, 'Start transition'
            print *, ''
                
            ! Transition
            r_TR = r_lr                                     ! interest rate path
            lambda_TR = lambda_lr                           ! tax level path
            mt_TR = mt_lr                                   ! transfer level path
            wge_TR = (1.0D0-alpha)/(r_TR+delta)             ! wage path
            wge_TR = alpha*(wge_TR**((1d0-alpha)/alpha))    ! wage path
        
            ! Compute Jacobian
            call Compute_JACOB

            ! Compute transition
            call Compute_QNEWTON
        
            
            ! Welfare accounting for transition to new steady state
            fvalSOB(ia) = -sum(mu_TR(:,1)*V_TR(:,1))     
            print *, 'Welfare transition = ', welfare
            print *, ''
        
            ! Store errors
            errors_vec_Sob(ia,:) = error_vec        ! steady state errors
            errors_vec_trans(ia,1,:) = asset_TR     ! transition: deviations from asset market clearing per period
            errors_vec_trans(ia,2,:) = BC_TR        ! transition: deviations from govt BC clearing per period
        
        ! If steady state did not converge: go to next point
        else      
            
            ! Reset eqm objects
            r =  r_init         ! interest rate  
            mt = 0.2d0          ! transfer level
            
            ! Wage implied by interest rate
            wge = (1.0D0-alpha)/(r+delta)
            wge = alpha*(wge**((1D0-alpha)/alpha))
            
            ! Very high value for (negative) welfare
            fvalSOB(ia) = 10000d0
        
            ! Store errors
            errors_vec_Sob(ia,:) = error_vec    ! steady state errors
            errors_vec_trans(ia,1,:) = 1d0      ! transition: deviations from asset market clearing per period
            errors_vec_trans(ia,2,:) = 1d0      ! transition: deviations from govt BC clearing per period
          
        endif
        
        ! Abs. maximum deviations from asset market and govt budget clearing along transition
        error_scalar_trans(ia,1) = maxval(abs(errors_vec_trans(ia,1,:)))  ! asset market 
        error_scalar_trans(ia,2) = maxval(abs(errors_vec_trans(ia,2,:)))  ! BC  
        
        ! Fill info variable
        Info(ia,:) = [ dble(ia), gamma, mt, rt,  error_scalar_trans(ia,1), error_scalar_trans(ia,2), fvalSOB_ss(ia),  fvalSOB(ia),    errors_vec_Sob(ia,3),   errors_vec_Sob(ia,4),  errors_vec_Sob(ia,1),  errors_vec_Sob(ia,2), lambda_lr, r_lr, dble(index_exit), dble(index_it) ]
        
        ! Write to file results from info variable
         inquire(file="Output/info.txt", exist=exist)
         if (exist) then
            open(145, file = "Output/info.txt", status="old", position="append", action="write")
         else
            open(145, file="Output/info.txt", status="new", action="write")
         end if
         write(145, '(*(f18.12))') ( Info(ia, isave), isave = 1, N_info)
         close(145)
         enddo
    
end program
